import os
import numpy as np
import acl
from PIL import Image

from constant import ACL_MEM_MALLOC_HUGE_FIRST, \
    ACL_MEMCPY_HOST_TO_DEVICE, ACL_MEMCPY_DEVICE_TO_HOST, \
    ACL_SUCCESS, NPY_FLOAT32

# Buffer methods for input and output sizes
buffer_method = {
    "in": acl.mdl.get_input_size_by_index,
    "out": acl.mdl.get_output_size_by_index
}

def check_ret(message, ret):
    """检查返回数据"""
    if ret != ACL_SUCCESS:
        raise Exception("{} failed ret={}".format(message, ret))

class Net(object):
    """推理网络定义"""
    def __init__(self, device_id, model_path, idx2label_list):
        self.device_id = device_id
        self.model_path = model_path
        self.model_id = None
        self.context = None

        self.input_data = []
        self.output_data = []
        self.model_desc = None
        self.load_input_dataset = None
        self.load_output_dataset = None

        self.init_resource()
        self.idx2label_list = idx2label_list

    def release_resource(self):
        """释放ACL资源"""
        print("Releasing resources stage:")
        ret = acl.mdl.unload(self.model_id)
        check_ret("acl.mdl.unload", ret)
        if self.model_desc:
            acl.mdl.destroy_desc(self.model_desc)
            self.model_desc = None

        while self.input_data:
            item = self.input_data.pop()
            ret = acl.rt.free(item["buffer"])
            check_ret("acl.rt.free", ret)

        while self.output_data:
            item = self.output_data.pop()
            ret = acl.rt.free(item["buffer"])
            check_ret("acl.rt.free", ret)

        if self.context:
            ret = acl.rt.destroy_context(self.context)
            check_ret("acl.rt.destroy_context", ret)
            self.context = None

        ret = acl.rt.reset_device(self.device_id)
        check_ret("acl.rt.reset_device", ret)
        print('Resources released successfully.')

    def init_resource(self):
        """初始化ACL资源"""
        acl.init()
        print("init resource stage:")
        ret = acl.rt.set_device(self.device_id)
        check_ret("acl.rt.set_device", ret)

        self.context, ret = acl.rt.create_context(self.device_id)
        check_ret("acl.rt.create_context", ret)

        # Load the model
        self.model_id, ret = acl.mdl.load_from_file(self.model_path)
        check_ret("acl.mdl.load_from_file", ret)
        print("model_id:{}".format(self.model_id))

        self.model_desc = acl.mdl.create_desc()
        self._get_model_info()
        print("init resource success")

    def _get_model_info(self):
        """获取模型信息"""
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        check_ret("acl.mdl.get_desc", ret)
        input_size = acl.mdl.get_num_inputs(self.model_desc)
        output_size = acl.mdl.get_num_outputs(self.model_desc)
        self._gen_data_buffer(input_size, des="in")
        self._gen_data_buffer(output_size, des="out")

    def _gen_data_buffer(self, size, des):
        """准备模型输入输出数据"""
        func = buffer_method[des]
        for i in range(size):
            temp_buffer_size = func(self.model_desc, i)
            temp_buffer, ret = acl.rt.malloc(temp_buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
            check_ret("acl.rt.malloc", ret)

            if des == "in":
                self.input_data.append({"buffer": temp_buffer, "size": temp_buffer_size})
            elif des == "out":
                self.output_data.append({"buffer": temp_buffer, "size": temp_buffer_size})

    def _data_interaction(self, dataset, policy=ACL_MEMCPY_HOST_TO_DEVICE):
        """在host和device之间通过复制内存进行数据传输"""
        temp_data_buffer = self.input_data if policy == ACL_MEMCPY_HOST_TO_DEVICE else self.output_data
        if len(dataset) == 0 and policy == ACL_MEMCPY_DEVICE_TO_HOST:
            for item in self.output_data:
                temp, ret = acl.rt.malloc_host(item["size"])
                if ret != 0:
                    raise Exception("can't malloc_host ret={}".format(ret))
                dataset.append({"size": item["size"], "buffer": temp})

        for i, item in enumerate(temp_data_buffer):
            if policy == ACL_MEMCPY_HOST_TO_DEVICE:
                # Ensure the data is contiguous and convert to bytes
                ptr = acl.util.bytes_to_ptr(np.ascontiguousarray(dataset[i]).tobytes())
                ret = acl.rt.memcpy(item["buffer"], item["size"], ptr, item["size"], policy)
                check_ret("acl.rt.memcpy", ret)
            else:
                ptr = dataset[i]["buffer"]
                ret = acl.rt.memcpy(ptr, item["size"], item["buffer"], item["size"], policy)
                check_ret("acl.rt.memcpy", ret)


    def _gen_dataset(self, type_str="input"):
        """准备模型输入输出数据集"""
        dataset = acl.mdl.create_dataset()
        temp_dataset = self.input_data if type_str == "in" else self.output_data

        for item in temp_dataset:
            data = acl.create_data_buffer(item["buffer"], item["size"])
            _, ret = acl.mdl.add_dataset_buffer(dataset, data)

            if ret != ACL_SUCCESS:
                ret = acl.destroy_data_buffer(data)
                check_ret("acl.destroy_data_buffer", ret)

        if type_str == "in":
            self.load_input_dataset = dataset
        else:
            self.load_output_dataset = dataset

    def _data_from_host_to_device(self, images):
        """将数据从host传到device"""
        print("data interaction from host to device")
        self._data_interaction(images, ACL_MEMCPY_HOST_TO_DEVICE)
        self._gen_dataset("in")
        self._gen_dataset("out")
        print("data interaction from host to device success")

    def _data_from_device_to_host(self):
        """将数据从device传到host"""
        print("data interaction from device to host")
        res = []
        self._data_interaction(res, ACL_MEMCPY_DEVICE_TO_HOST)
        print("data interaction from device to host success")
        result = self.get_result(res)
        pred_dict = self._print_result(result)
        for item in res:
            ptr = item['buffer']
            ret = acl.rt.free_host(ptr)
            check_ret('acl.rt.free_host', ret)
        return pred_dict

    def run(self, images):
        """执行数据传输和模型推理"""
        self._data_from_host_to_device(images)
        self.forward()
        pred_dict = self._data_from_device_to_host()
        return pred_dict

    def forward(self):
        """模型推理"""
        print('execute stage:')
        ret = acl.mdl.execute(self.model_id, self.load_input_dataset, self.load_output_dataset)
        check_ret("acl.mdl.execute", ret)
        self._destroy_databuffer()
        print('execute stage success')

    def _print_result(self, result):
        """打印推理结果"""
        vals = np.array(result).flatten()
        top_k = vals.argsort()[-1:-6:-1]
        print("======== top5 inference results: =============")
        
        pred_dict = {}
        for j in top_k:
            print(f'{self.idx2label_list[j]}: {vals[j]}')
            pred_dict[self.idx2label_list[j]] = vals[j]
        return pred_dict

    def _destroy_databuffer(self):
        """释放模型输入输出数据"""
        for dataset in [self.load_input_dataset, self.load_output_dataset]:
            if not dataset:
                continue
            number = acl.mdl.get_dataset_num_buffers(dataset)
            for i in range(number):
                data_buf = acl.mdl.get_dataset_buffer(dataset, i)
                if data_buf:
                    ret = acl.destroy_data_buffer(data_buf)
                    check_ret("acl.destroy_data_buffer", ret)
            ret = acl.mdl.destroy_dataset(dataset)
            check_ret("acl.mdl.destroy_dataset", ret)

    def get_result(self, output_data):
        """获取模型输出结果"""
        result = []
        dims, ret = acl.mdl.get_cur_output_dims(self.model_desc, 0)
        check_ret("acl.mdl.get_cur_output_dims", ret)
        out_dim = dims['dims']
        for temp in output_data:
            ptr = temp["buffer"]
            data = acl.util.ptr_to_numpy(ptr, tuple(out_dim), NPY_FLOAT32)
            result.append(data)
        return result


def preprocess_img(input_path):
    """图片预处理"""
    input_path = os.path.abspath(input_path)
    with Image.open(input_path) as image_file:
        image_file = image_file.resize((224, 224))  # Resize to 224x224
        img = np.array(image_file)
    img = img.astype("float32")  # Change to float32 for MobileNetV3
    # Normalize using MobileNetV3 specific mean and std
    img[:, :, 0] = (img[:, :, 0] / 255.0 - 0.485) / 0.229  # Normalize R channel
    img[:, :, 1] = (img[:, :, 1] / 255.0 - 0.456) / 0.224  # Normalize G channel
    img[:, :, 2] = (img[:, :, 2] / 255.0 - 0.406) / 0.225  # Normalize B channel
    img = img.transpose((2, 0, 1))  # 转换为CHW格式
    return img


def load_idx2label(label_path):
    """加载标签映射"""
    with open(label_path, 'r') as f:
        idx2label = [line.strip() for line in f.readlines()]
    return idx2label


def main(input_image_path, model_path, label_path):
    """主程序入口"""
    device_id = 0
    idx2label_list = load_idx2label(label_path)
    net = Net(device_id, model_path, idx2label_list)
    img = preprocess_img(input_image_path)
    pred_dict = net.run([img])
    net.release_resource()


if __name__ == "__main__":
    input_image_path = '22.png'  # 你的输入图片路径
    model_path = 'model.om'  # 你的模型路径
    label_path = 'names.txt'  # 你的标签路径，修改为TXT文件
    main(input_image_path, model_path, label_path)
